import aiohttp
import asyncio
import base64
import os
import json
import time
from typing import List, Optional
import logging
from config.setting import VOLCANO_API_KEY, VOLCANO_API_URL, VOLCANO_MODEL_NAME

logger = logging.getLogger(__name__)


class VolcanoClient:
    def __init__(self):
        self.api_key = VOLCANO_API_KEY
        self.api_url = VOLCANO_API_URL
        self.model_name = VOLCANO_MODEL_NAME
        self.supported_image_formats = {
            '.jpg': 'jpeg', '.jpeg': 'jpeg', '.png': 'png', '.gif': 'gif',
            '.webp': 'webp', '.bmp': 'bmp', '.tiff': 'tiff', '.ico': 'ico',
            '.heic': 'heic', '.heif': 'heif'
        }
        self.client = aiohttp.ClientSession()

    async def close(self):
        """关闭HTTP会话"""
        await self.client.close()

    def encode_image_to_base64(self, image_path):
        """将图片编码为base64格式"""
        try:
            ext = os.path.splitext(image_path.lower())[1]
            if ext not in self.supported_image_formats:
                logger.warning(f"不支持的图片格式: {ext}")
                return None, None

            img_format = self.supported_image_formats[ext]
            with open(image_path, "rb") as image_file:
                image_data = base64.b64encode(image_file.read()).decode('utf-8')
            return img_format, image_data
        except Exception as e:
            logger.error(f"图片编码失败: {str(e)}")
            return None, None

    async def call_volcano_api(
            self,
            system_prompt: str,
            user_prompt: str,
            image_paths: Optional[List[str]] = None,
            max_retries: int = 3
    ) -> str:
        """调用火山API"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        # 构建消息体
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": [{"type": "text", "text": user_prompt}]}
        ]

        # 添加图片
        if image_paths:
            for img_path in image_paths[:3]:  # 最多3张图
                img_format, img_base64 = self.encode_image_to_base64(img_path)
                if img_format and img_base64:
                    messages[1]["content"].append({
                        "type": "image_url",
                        "image_url": {"url": f"data:image/{img_format};base64,{img_base64}"}
                    })

        payload = {
            "model": self.model_name,
            "messages": messages,
            "temperature": 0.1,
            "max_tokens": 2048,
            "stream": False,
            "thinking": {"type": "disabled"}
        }

        # 带重试的请求
        for attempt in range(max_retries):
            try:
                async with self.client.post(
                        self.api_url,
                        headers=headers,
                        json=payload,
                        timeout=60
                ) as response:
                    response_text = await response.text()
                    logger.info(f"火山API响应: {response_text[:500]}...")

                    if response.status in [400, 401, 404]:
                        return f"API错误 {response.status}: {response_text}"

                    response.raise_for_status()
                    result = await response.json()
                    if "choices" in result and result["choices"]:
                        return result["choices"][0]["message"]["content"]
                    return f"空结果: {result}"

            except Exception as e:
                logger.error(f"API调用失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt < max_retries - 1:
                    await asyncio.sleep(5 * (attempt + 1))

        return "超过最大重试次数"


async def get_volcano_client():
    """依赖注入：获取火山客户端实例"""
    client = VolcanoClient()
    try:
        yield client
    finally:
        await client.close()
